from dml import TripletModel,TripletDataset,X_train,y_train,triplet_loss
import torch
from torch.utils.data import DataLoader, Dataset


# 加载训练好的模型
model = TripletModel()
model.load_state_dict(torch.load('src/model_pth/triplet_model_iu_xray.pth',weights_only=False))
model.eval()

# 创建测试数据集
num_triplets = 100
test_dataset = TripletDataset(X_train, y_train, num_triplets)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

# 测试模型
correct = 0
total = 0
with torch.no_grad():
    for batch in test_loader:
        anchor, positive, negative = batch

        anchor = anchor.unsqueeze(1)  # 增加通道维度
        positive = positive.unsqueeze(1)  # 增加通道维度
        negative = negative.unsqueeze(1)  # 增加通道维度

        anchor_embedding = model(anchor)
        positive_embedding = model(positive)
        negative_embedding = model(negative)
        
        # 计算三元组损失
        loss = triplet_loss(anchor_embedding, positive_embedding, negative_embedding)
        
        # 计算准确率
        pos_dist = torch.sum(torch.square(anchor_embedding - positive_embedding), dim=1)
        neg_dist = torch.sum(torch.square(anchor_embedding - negative_embedding), dim=1)
        correct += torch.sum(pos_dist < neg_dist).item()
        total += anchor.size(0)

accuracy = correct / total
print(f'Test Accuracy: {accuracy * 100:.2f}%')